import uuid
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
import time
from typing import Generator, Dict
import requests
import json
import logging
import sqlite3
import threading
app = FastAPI()
logger = logging.getLogger(__name__)
# 配置logging模块
handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
lock = threading.Lock()

def insert_to_db(request_id: str, data: dict, ret_data: str = None):
    """Insert or update log data into the database"""
    with lock:
        conn = sqlite3.connect('logs.db')  # Create a new connection inside the function
        cursor = conn.cursor()
        
        # Check if there is an existing record with the same request_id
        cursor.execute("SELECT * FROM logs WHERE request_id=?", (request_id,))
        existing_record = cursor.fetchone()
        
        if existing_record:
            # If a record exists, update it
            cursor.execute('''UPDATE logs SET data=?, ret_data=? WHERE request_id=?''', (json.dumps(data, ensure_ascii=False), ret_data, request_id))
        else:
            # Otherwise, insert a new record
            if ret_data is not None:
                cursor.execute('''INSERT INTO logs (request_id, data, ret_data) VALUES (?, ?, ?)''', (request_id, json.dumps(data, ensure_ascii=False), ret_data)) 
            else:
                cursor.execute('''INSERT INTO logs (request_id, data) VALUES (?, ?)''', (request_id, json.dumps(data, ensure_ascii=False)))
        
        conn.commit()
        conn.close()  # Close the connection after use

def parse_all(data):
    role = None
    contents = []
    for chunk in data:
        if not (chunk.startswith('data:') and '[DONE]' not in chunk):
            continue
        try:
            chunk_data = json.loads(chunk[5:])
            delta = chunk_data['choices'][0]['delta']
            contents.append(delta['content'])
            if role is None:
                role = delta['role']
        except (IndexError, KeyError, ValueError):
            print(f"Warning: Invalid JSON format in line '{chunk}'")
    return json.dumps({'role': role, 'content': ''.join(contents)}, ensure_ascii=False)

def generate_sse_events(url: str, data: Dict) -> Generator[str, None, None]:
    """从指定URL获取事件流并转发"""
    headers = {"content-type": "application/json"}
    try:
        request_id=str(uuid.uuid4())
        insert_to_db(request_id, data)  # write data to the logs table
        logger.info(f"request id {request_id} request: {data}")
        response = requests.post(url, json=data, headers=headers, stream=True, timeout=30)
        if response.status_code != 200:
            yield f"Event: error\nData: Request failed with status code {response.status_code}\n\n"
            return
    except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
        yield f"Event: error\nData: Request failed: {str(e)}\n\n"
        return
    response_data=[]
    for line in response.iter_lines():
        if line:
            response_data.append(line.decode('utf-8'))
            yield line.decode('utf-8')+'\n\n'
    # write ret_data to the logs table
    ret_data=parse_all(response_data)
    insert_to_db(request_id, data, ret_data)
    logger.info(f"request id {request_id} response: {ret_data}",)
    yield f"Event: response\nData: {ret_data}\n\n"

@app.post("/v1/chat/completions")
async def sse_events(request: Request):
    """SSE端点，从另外一个网站获取事件流并转发给用户"""
    url = "https://api.deepseek.com/v1/chat/completions"  # 替换为需要获取事件流的URL
    data = await request.json()
    event_generator = generate_sse_events(url, data)
    return StreamingResponse(event_generator, media_type="text/event-stream", headers={"X-Request-ID": str(uuid.uuid4())})

if __name__ == "__main__":
    import uvicorn

    # Create a SQLite database connection and cursor
    conn = sqlite3.connect('logs.db')
    cursor = conn.cursor()
    # Create the logs table if it doesn't exist
    cursor.execute('''CREATE TABLE IF NOT EXISTS logs (request_id TEXT PRIMARY KEY, data TEXT, ret_data TEXT)''')
    conn.commit()
    conn.close()

    uvicorn.run(app, host="0.0.0.0", port=8200)
